import sys
import cv2
import csv
import datetime
import os
import numpy as np
from PyQt5.QtWidgets import *
from PyQt5.QtGui import *
from PyQt5.QtCore import *


class FaceCaptureApp(QWidget):
    def __init__(self):
        super().__init__()
        self.cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        self.face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
        self.face_recognizer = cv2.face.LBPHFaceRecognizer_create()

        self.attendance_records = []
        self.current_session = datetime.datetime.now().strftime("%Y%m%d")
        self.record_file = f"attendance_{self.current_session}.csv"
        self.load_attendance()

        self.registered_students = {}

        self.frame_counter = 0
        self.is_paused = False
        self.registration_mode = False
        self.current_register_id = None
        self.current_register_name = None
        self.registration_counter = 0
        self.export_path = "exports"

        # 先初始化UI，创建status_label
        self.init_ui()

        # 再加载注册学生信息
        self.load_registered_students()

        self.timer = self.startTimer(30)

    def init_ui(self):
        main_layout = QVBoxLayout()

        # 视频显示区域
        video_layout = QHBoxLayout()
        self.video_label = QLabel()
        video_layout.addWidget(self.video_label)

        # 右侧控制面板
        control_layout = QVBoxLayout()

        # 输入区域、按钮等布局...

        # 考勤记录显示
        self.record_table = QTableWidget()
        control_layout.addWidget(self.record_table)

        # 整合布局
        video_layout.addLayout(control_layout)
        main_layout.addLayout(video_layout)

        self.setLayout(main_layout)
        self.setWindowTitle("智能人脸识别考勤系统")
        self.setGeometry(100, 100, 800, 600)

        # 创建并添加status_label到布局
        self.status_label = QLabel()
        main_layout.addWidget(self.status_label)

        # 其他方法保持不变...
        # 输入区域（仅用于手动签到）
        input_group = QGroupBox("学生信息（手动签到时使用）")
        input_layout = QFormLayout()
        self.name_edit = QLineEdit(placeholderText="姓名（自动填充）")
        self.id_edit = QLineEdit(placeholderText="学号（自动填充）")
        input_layout.addRow("姓名:", self.name_edit)
        input_layout.addRow("学号:", self.id_edit)
        input_group.setLayout(input_layout)
        control_layout.addWidget(input_group)

        # 控制按钮
        button_layout = QHBoxLayout()
        self.play_pause_button = QPushButton("暂停")
        self.play_pause_button.clicked.connect(self.toggle_play_pause)
        self.register_button = QPushButton("注册学生")
        self.register_button.clicked.connect(self.register_student)
        self.capture_button = QPushButton("人脸识别签到")
        self.capture_button.clicked.connect(self.face_recognition_sign_in)
        self.export_button = QPushButton("导出记录")
        self.export_button.clicked.connect(self.export_records)
        self.path_button = QPushButton("设置导出路径")
        self.path_button.clicked.connect(self.set_export_path)
        button_layout.addWidget(self.play_pause_button)
        button_layout.addWidget(self.register_button)
        button_layout.addWidget(self.capture_button)
        button_layout.addWidget(self.path_button)
        control_layout.addLayout(button_layout)

        # 考勤记录显示
        self.record_table = QTableWidget()
        self.record_table.setColumnCount(4)
        self.record_table.setHorizontalHeaderLabels(["序号", "学号", "姓名", "时间"])
        self.record_table.setEditTriggers(QAbstractItemView.NoEditTriggers)
        control_layout.addWidget(QLabel("今日考勤记录:"))
        control_layout.addWidget(self.record_table)

        # 状态提示
        self.status_label = QLabel()  # 初始化 status_label
        main_layout.addWidget(self.status_label)

        # 整合布局
        video_layout.addLayout(control_layout)
        main_layout.addLayout(video_layout)

        self.setLayout(main_layout)
        self.setWindowTitle("智能人脸识别考勤系统")
        self.setGeometry(100, 100, 800, 600)

    def timerEvent(self, event):
        self.update_frame()

    def update_frame(self):
        if not self.is_paused:
            ret, frame = self.cap.read()
            if not ret:
                self.show_status("摄像头读取失败，请检查设备", 3000)
                return

            self.frame_counter += 1
            if self.frame_counter % 3 != 0:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                q_image = QImage(frame.data, 640, 480, 3 * 640, QImage.Format_RGB888)
                self.video_label.setPixmap(QPixmap.fromImage(q_image))
                return

            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)

            for (x, y, w, h) in faces:
                cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)

            # 注册模式下保存图像
            if self.registration_mode:
                if len(faces) == 1:
                    (x, y, w, h) = faces[0]
                    face_gray = gray[y:y + h, x:x + w]
                    filename = f"registered_faces/{self.current_register_id}/{self.registration_counter}.jpg"
                    os.makedirs(os.path.dirname(filename), exist_ok=True)
                    cv2.imwrite(filename, face_gray)
                    self.registration_counter += 1
                    if self.registration_counter >= 50:
                        self.registration_mode = False
                        self.show_status(f"注册成功！学号：{self.current_register_id}", 3000)
                        self.load_registered_students()

            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            q_image = QImage(frame.data, 640, 480, 3 * 640, QImage.Format_RGB888)
            self.video_label.setPixmap(QPixmap.fromImage(q_image))

    def toggle_play_pause(self):
        self.is_paused = not self.is_paused
        self.play_pause_button.setText("播放" if self.is_paused else "暂停")

    def register_student(self):
        name, ok = QInputDialog.getText(self, "注册学生", "姓名:")
        if not ok or not name:
            self.show_status("注册失败：请输入姓名", 3000)
            return
        id_num, ok = QInputDialog.getText(self, "注册学生", "学号（数字）:")
        if not ok or not id_num:
            self.show_status("注册失败：请输入学号", 3000)
            return
        try:
            id_num_int = int(id_num)
        except ValueError:
            self.show_status("学号必须为整数", 3000)
            return
        if str(id_num_int) in self.registered_students:
            self.show_status("学号已存在，请更换学号", 3000)
            return

        # 保存学生信息到CSV
        with open('registered_students.csv', 'a', newline='') as f:
            writer = csv.writer(f)
            writer.writerow([id_num, name])

        self.current_register_id = id_num
        self.current_register_name = name
        self.registration_counter = 0
        self.registration_mode = True
        os.makedirs(f"registered_faces/{id_num}", exist_ok=True)
        self.show_status("开始采集人脸，请保持面部在框内...", 5000)

    def face_recognition_sign_in(self):
        ret, frame = self.cap.read()
        if not ret:
            self.show_status("摄像头读取失败，请检查设备", 3000)
            return
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)
        if len(faces) != 1:
            self.show_status("检测到多个人脸或未检测到人脸，请调整位置", 3000)
            return
        (x, y, w, h) = faces[0]
        face_gray = gray[y:y + h, x:x + w]

        try:
            id_pred, confidence = self.face_recognizer.predict(face_gray)
        except cv2.error:
            self.show_status("未注册学生，请先注册", 3000)
            return

        if confidence < 80:
            student_id = str(id_pred)
            student = self.registered_students.get(student_id, None)
            if student:
                record_time = datetime.datetime.now().strftime("%H:%M:%S")
                new_record = [
                    len(self.attendance_records) + 1,
                    student_id,
                    student['name'],
                    record_time
                ]
                self.attendance_records.append(new_record)
                self.update_record_table()
                self.save_attendance()
                self.show_status(f"签到成功！学号：{student_id}", 3000)
            else:
                self.show_status("未找到该学生信息，请重新注册", 3000)
        else:
            self.show_status("人脸识别失败，请手动输入信息进行签到", 3000)

    def load_attendance(self):
        if os.path.exists(self.record_file):
            with open(self.record_file, 'r', newline='', encoding='utf-8') as f:
                reader = csv.reader(f)
                next(reader)  # 跳过标题行
                self.attendance_records = [row for row in reader]

    def save_attendance(self):
        try:
            with open(self.record_file, 'w', newline='', encoding='utf-8') as f:
                writer = csv.writer(f)
                writer.writerow(["序号", "学号", "姓名", "时间"])
                writer.writerows(self.attendance_records)
        except Exception as e:
            self.show_status(f"保存考勤记录失败：{str(e)}", 3000)

    def load_registered_students(self):
        self.registered_students.clear()
        if os.path.exists('registered_students.csv'):
            with open('registered_students.csv', 'r') as f:
                reader = csv.reader(f)
                for row in reader:
                    if len(row) >= 2:
                        student_id = row[0]
                        name = row[1]
                        self.registered_students[student_id] = {'id': student_id, 'name': name}

        # 加载图像和训练模型
        faces_dir = r"D:\人脸识别考勤系统\记录"
        images = []
        labels = []
        for student_id in os.listdir(faces_dir):
            if not student_id.isdigit():
                continue
            student_id_int = int(student_id)
            student_path = os.path.join(faces_dir, student_id)
            if not os.path.isdir(student_path):
                continue
            for img_name in os.listdir(student_path):
                img_path = os.path.join(student_path, img_name)
                img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
                if img is None:
                    continue
                images.append(img)
                labels.append(student_id_int)
        if len(images) > 0:
            self.face_recognizer.train(images, np.array(labels))
        else:
            self.show_status("未找到注册学生数据", 3000)

    def update_record_table(self):
        self.record_table.setRowCount(len(self.attendance_records))
        for row, record in enumerate(self.attendance_records):
            for col, item in enumerate(record):
                self.record_table.setItem(row, col, QTableWidgetItem(str(item)))

    def set_export_path(self):
        path = QFileDialog.getExistingDirectory(self, "选择导出文件夹")
        if path:
            self.export_path = path
            self.show_status(f"导出路径设置为：{path}", 3000)

    def export_records(self):
        if not hasattr(self, 'export_path') or not self.export_path:
            self.show_status("请先设置导出路径", 3000)
            return
        filename = os.path.join(self.export_path, f"attendance_{self.current_session}.csv")
        try:
            self.save_attendance_as_csv(filename)
        except Exception as e:
            self.show_status(f"导出失败：{str(e)}", 3000)

    def save_attendance_as_csv(self, filename):
        try:
            with open(filename, 'w', newline='', encoding='utf-8') as f:
                writer = csv.writer(f)
                writer.writerow(["序号", "学号", "姓名", "时间"])
                writer.writerows(self.attendance_records)
            self.show_status(f"导出成功！文件路径：{filename}", 3000)
        except Exception as e:
            self.show_status(f"导出CSV失败：{str(e)}", 3000)

    def show_status(self, message, duration=2000):
        self.status_label.setText(message)
        QTimer.singleShot(duration, self.status_label.clear)

    def closeEvent(self, event):
        self.cap.release()
        event.accept()

if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = FaceCaptureApp()
    ex.show()
    sys.exit(app.exec_())
